// RemoveReadOnlyContextMenuExt.cpp
// (c) 2008, Charles Lechasseur
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

#include "stdafx.h"
#include "RemoveReadOnlyContextMenuExt.h"

#include "resource.h"
#include "StStgMedium.h"
#include "StCursor.h"


namespace {

    // Maximum number of files that can be selected if we want to pre-compute a flag state.
    const size_t MAX_NUM_FILES_FOR_PRECOMPUTE = 16;

}; // anonymous namespace


// CRemoveReadOnlyContextMenuExt

//
// IShellExtInit::Initialize
//
// Called by the shell to initialize our contextual menu extension.
// We need to use this opportunity to look at the files/folders to act upon.
//
// @param p_pFolderPIDL Pointer to ITEMIDLIST representing selected folder; unused.
// @param p_pDataObject Pointer to data object containing information about selected files/folders.
// @param p_hKeyFileClass Handle to file class key; unused.
// @return S_OK if successful, otherwise an error code.
//
STDMETHODIMP CRemoveReadOnlyContextMenuExt::Initialize(
    PCIDLIST_ABSOLUTE /*p_pFolderPIDL*/,
    IDataObject* p_pDataObject,
    HKEY /*p_hKeyFileClass*/)
{
    HRESULT hRes = S_OK;

    try {
        // Make sure we have a data object.
        if (p_pDataObject == 0) {
            hRes = E_INVALIDARG;
        } else {
            // Extract HDROP from data object.
            StStgMedium stgMedium;
            FORMATETC formatEtc = {CF_HDROP, 0, DVASPECT_CONTENT, -1, TYMED_HGLOBAL};
            if (FAILED(p_pDataObject->GetData(&formatEtc, &stgMedium))) {
                // It's difficult to display a menu item if OS doesn't work.
                hRes = E_FAIL;
            } else {
                // Get number of files included in the selection.
                UINT fileCount = ::DragQueryFileW(
                    static_cast<HDROP>(stgMedium.get().hGlobal), 0xFFFFFFFF, 0, 0);
                if (fileCount == 0) {
                    // It's difficult to display a menu item without files to act upon.
                    hRes = E_FAIL;
                } else {
                    // Pre-allocate space in vector to store files.
                    m_vFiles.reserve(fileCount);

                    // Get each file in turn.
                    wchar_t buffer[MAX_PATH + 1];
                    for(UINT i = 0; i < fileCount; ++i) {
                        UINT copiedCount = ::DragQueryFileW(static_cast<HDROP>(stgMedium.get().hGlobal),
                            i, buffer, sizeof(buffer) / sizeof(wchar_t));
                        m_vFiles.push_back(std::wstring(buffer, copiedCount));
                    }

                    // Fun "optimization": if user did not select too many files and did
                    // not select a directory, we will try to "precompute" the flag state
                    // that he/she is interested in (e.g., if all files are already read-only,
                    // we will only show the Remove command).
                    if (m_vFiles.size() <= MAX_NUM_FILES_FOR_PRECOMPUTE) {
                        FilesV::const_iterator it, end = m_vFiles.end();
                        for (it = m_vFiles.begin(); it != end; ++it) {
                            DWORD attributes = ::GetFileAttributesW(it->c_str());
                            if (attributes == INVALID_FILE_ATTRIBUTES) {
                                // Error, can't precompute.
                                m_PredefinedFlagState.ClearValue();
                                break;
                            } else {
                                bool readOnly = (attributes & FILE_ATTRIBUTE_READONLY) != 0;
                                if (((attributes & FILE_ATTRIBUTE_DIRECTORY) != 0)
                                    || (m_PredefinedFlagState.HasValue()
                                        && (m_PredefinedFlagState != !readOnly))) {

                                    // Either a directory or a different flag, can't precompute.
                                    m_PredefinedFlagState.ClearValue();
                                    break;
                                } else {
                                    m_PredefinedFlagState = !readOnly;
                                }
                            }
                        }
                    }
                }
            }
        }
    } catch (...) {
        hRes = E_UNEXPECTED;
    }

    return hRes;
}

//
// IContextMenu::QueryContextMenu
//
// Invoked by the shell to populate a contextual menu.
// We need to use this opportunity to add our menu items.
//
// @param p_hMenu Handle to the contextual menu to populate.
// @param p_Index Index in the menu where to add items.
// @param p_FirstCmdId First available command ID for our commands.
// @param p_LastCmdId Last available command ID for our commands.
// @param p_Flags Flags indicating how the menu can be changed; see MSDN for details.
// @return If successful, a success code with code value set to the largest command ID
//         used, plus one; otherwise, an error code.
//
STDMETHODIMP CRemoveReadOnlyContextMenuExt::QueryContextMenu(
    HMENU p_hMenu,
    UINT p_Index,
    UINT p_FirstCmdId,
    UINT p_LastCmdId,
    UINT p_Flags)
{
    HRESULT hRes = S_OK;

    try {
        if (p_hMenu == 0) {
            hRes = E_INVALIDARG;
        } else {
            // Do not add items if the default action is chosen, or if we have no files.
            if (!m_vFiles.empty() && ((p_Flags & CMF_DEFAULTONLY) == 0)) {
                UINT position = p_Index;
                UINT cmdId = p_FirstCmdId;

                // Add "Remove Read-only" menu item (unless predefined state is "Set").
                if ((cmdId <= p_LastCmdId)
                    && (!m_PredefinedFlagState.HasValue() || !m_PredefinedFlagState)) {

                    CString rroCaption(MAKEINTRESOURCE(IDS_REMOVE_READ_ONLY_MENU_ITEM));
#ifdef _DEBUG
                    rroCaption += _T(" (DEBUG)");
#endif
                    if (::InsertMenu(p_hMenu, position, MF_STRING | MF_BYPOSITION, cmdId, rroCaption)) {
                        // It worked, save command ID and go to next position/ID.
                        m_RemoveCmdId = cmdId;
                        if (!m_FirstCmdId.HasValue()) {
                            m_FirstCmdId = cmdId;
                        }
                        ++position;
                        ++cmdId;
                    } else {
                        hRes = E_FAIL;
                    }
                }

                // Add "Set Read-only" menu item (unless predefined state is "Remove").
                if (SUCCEEDED(hRes) && (cmdId <= p_LastCmdId)
                    && (!m_PredefinedFlagState.HasValue() || m_PredefinedFlagState)) {

                    CString sroCaption(MAKEINTRESOURCE(IDS_SET_READ_ONLY_MENU_ITEM));
#ifdef _DEBUG
                    sroCaption += _T(" (DEBUG)");
#endif
                    if (::InsertMenu(p_hMenu, position, MF_STRING | MF_BYPOSITION, cmdId, sroCaption)) {
                        // It worked, save command ID and go to next position/ID.
                        m_SetCmdId = cmdId;
                        if (!m_FirstCmdId.HasValue()) {
                            m_FirstCmdId = cmdId;
                        }
                        ++position;
                        ++cmdId;
                    } else {
                        hRes = E_FAIL;
                    }
                }

                if (SUCCEEDED(hRes)) {
                    // Strange return value requirement... see MSDN for details.
                    hRes = MAKE_HRESULT(SEVERITY_SUCCESS, 0, cmdId - p_FirstCmdId + 1);
                }
            }
        }
    } catch (...) {
        hRes = E_UNEXPECTED;
    }

    return hRes;
}

//
// IContextMenu::InvokeCommand
//
// Invoked by the shell when the user selects our menu item.
// We do our stuff here.
//
// @param p_pCommandInfo Pointer to struct containing command information.
// @return S_OK if successful, otherwise an error code.
//
STDMETHODIMP CRemoveReadOnlyContextMenuExt::InvokeCommand(
    CMINVOKECOMMANDINFO* p_pCommandInfo)
{
    HRESULT hRes = S_OK;

    try {
        if ((p_pCommandInfo == 0) || (p_pCommandInfo->cbSize < sizeof(CMINVOKECOMMANDINFO))) {
            hRes = E_INVALIDARG;
        } else {
            // Get offset of invoked command.
            UINT cmdOffset = (UINT) p_pCommandInfo->lpVerb;
            if ((cmdOffset & 0xFFFF0000) != 0) {
                // We do not support verb invokation.
                hRes = E_FAIL;
            } else {
                // Check which command it is that is invoked.
                Nullable<bool> isSro;
                if (m_FirstCmdId.HasValue()) {
                    if (m_RemoveCmdId.HasValue() && (cmdOffset == (m_RemoveCmdId - m_FirstCmdId))) {
                        isSro = false;
                    } else if (m_SetCmdId.HasValue() && (cmdOffset == (m_SetCmdId - m_FirstCmdId))) {
                        isSro = true;
                    }
                }
                if (!isSro.HasValue()) {
                    // This is not a recognized command ID.
                    hRes = E_FAIL;
                } else {
                    // We need to set or remove read-only flag on our files.
                    // Since this might take a while, change the cursor.
                    StCursor waitCursor(0, IDC_WAIT);

                    // Loop through files and modify flags.
                    FilesV::const_iterator it, end = m_vFiles.end();
                    for (it = m_vFiles.begin(); it != end; ++it) {
                        try {
                            ModifyReadOnlyFlag(*it, isSro);
                        } catch (...) {
                            // We want to try them all even if one of them fails.
                        }
                    }
                }
            }
        }
    } catch (...) {
        hRes = E_UNEXPECTED;
    }

    return hRes;
}

//
// IContextMenu::GetCommandString
//
// Invoked by the shell to get a text description for our menu item.
//
// @param p_CmdId Offset of ID of command for which to query text, relative to
//                our first command ID.
// @param p_Flags Type of information or action requested; see MSDN for details.
// @param p_pReserved Reserved; unused.
// @param p_pBuffer Pointer to memory buffer where to copy a null-terminated string.
//                  Note: if GCS_UNICODE is found in p_StringType, the string must
//                  be cast to a LPWSTR and a Unicode string must be stored there.
// @param p_BufferSize Max size of p_pBuffer, in characters. 
// @return S_OK if successful, otherwise an error code.
//
STDMETHODIMP CRemoveReadOnlyContextMenuExt::GetCommandString(
    UINT_PTR p_CmdId,
    UINT p_Flags,
    UINT* /*p_pReserved*/,
    LPSTR p_pBuffer,
    UINT p_BufferSize)
{
    HRESULT hRes = S_OK;

    try {
        // Check what is requested.
        if ((p_Flags == GCS_VERBA) || (p_Flags == GCS_VERBW)) {
            // We do not support verb invokation.
            hRes = E_NOTIMPL;
        } else if ((p_Flags == GCS_VALIDATEA) || (p_Flags == GCS_VALIDATEW)) {
            // We need to validate command ID.
            if (m_FirstCmdId.HasValue()) {
                bool isRro = m_RemoveCmdId.HasValue() ? (p_CmdId == (m_RemoveCmdId - m_FirstCmdId)) : false;
                bool isSro = m_SetCmdId.HasValue() ? (p_CmdId == (m_SetCmdId - m_FirstCmdId)) : false;
                if (!isRro && !isSro) {
                    // Invalid command ID.
                    hRes = S_FALSE;
                }
            } else {
                // No command IDs, so it's invalid.
                hRes = S_FALSE;
            }
        } else if (p_Flags == GCS_HELPTEXTA) {
            // An ANSI help string is requested.
            std::auto_ptr<CStringA> helpString;
            if (p_pBuffer == 0) {
                hRes = E_INVALIDARG;
            } else if (m_RemoveCmdId.HasValue() && (p_CmdId == (m_RemoveCmdId - m_FirstCmdId))) {
                helpString.reset(new CStringA(MAKEINTRESOURCEA(IDS_REMOVE_READ_ONLY_HINT)));
            } else if (m_SetCmdId.HasValue() && (p_CmdId == (m_SetCmdId - m_FirstCmdId))) {
                helpString.reset(new CStringA(MAKEINTRESOURCEA(IDS_SET_READ_ONLY_HINT)));
            } else {
                // Unknown command ID.
                hRes = E_INVALIDARG;
            }
            if (SUCCEEDED(hRes)) {
                if ((helpString.get() == 0)
                        || (static_cast<UINT>(helpString->GetLength()) >= p_BufferSize)) {
                    hRes = E_FAIL;
                } else {
                    ::strcpy(p_pBuffer, *helpString);
                }
            }
        } else if (p_Flags == GCS_HELPTEXTW) {
            // A Unicode help string is requested.
            std::auto_ptr<CStringW> helpString;
            if (p_pBuffer == 0) {
                hRes = E_INVALIDARG;
            } else if (m_RemoveCmdId.HasValue() && (p_CmdId == (m_RemoveCmdId - m_FirstCmdId))) {
                helpString.reset(new CStringW(MAKEINTRESOURCEW(IDS_REMOVE_READ_ONLY_HINT)));
            } else if (m_SetCmdId.HasValue() && (p_CmdId == (m_SetCmdId - m_FirstCmdId))) {
                helpString.reset(new CStringW(MAKEINTRESOURCEW(IDS_SET_READ_ONLY_HINT)));
            } else {
                // Unknown command ID.
                hRes = E_INVALIDARG;
            }
            if (SUCCEEDED(hRes)) {
                if ((helpString.get() == 0)
                        || (static_cast<UINT>(helpString->GetLength()) >= p_BufferSize)) {
                    hRes = E_FAIL;
                } else {
                    ::wcscpy(reinterpret_cast<LPWSTR>(p_pBuffer), *helpString);
                }
            }
        } else {
            // Unknown, unsupported flag.
            hRes = E_NOTIMPL;
        }
    } catch (...) {
        hRes = E_UNEXPECTED;
    }

    return hRes;
}

//
// Modifies the "read-only" flag on the given file or directory.
// If path points to a directory, the operation is recursive.
//
// @param p_Path Path of file or directory to change.
// @param p_FlagState true if we want to set the "read-only" flag,
//                    false if we want to remove it.
//
void
CRemoveReadOnlyContextMenuExt::ModifyReadOnlyFlag(
    const std::wstring& p_Path,
    const bool p_FlagState)
{
    // Get path attributes.
    DWORD attributes = ::GetFileAttributesW(p_Path.c_str());
    if (attributes != INVALID_FILE_ATTRIBUTES) {
        // Check if it's a file or directory.
        if ((attributes & FILE_ATTRIBUTE_DIRECTORY) == 0) {
            // This is a file.
            // Set or clear read-only flag and re-apply attributes.
            if (p_FlagState) {
                attributes |= FILE_ATTRIBUTE_READONLY;
            } else {
                attributes &= ~FILE_ATTRIBUTE_READONLY;
            }
            ::SetFileAttributesW(p_Path.c_str(), attributes);
        } else {
            // This is a directory.
            // Recurse into it.
            // (Note: since this is recursive, try to be minimalistic with stack space.)
            std::auto_ptr<FilesV> apSubPaths(new FilesV);
            std::auto_ptr<WIN32_FIND_DATAW> apFindData(new WIN32_FIND_DATAW);
            HANDLE hFind = ::FindFirstFileW((p_Path + L"\\*.*").c_str(), apFindData.get());
            if (hFind != INVALID_HANDLE_VALUE) {
                try {
                    do {
                        // Filter out "." and ".." (stupid Windows)
                        if ((::wcscmp(L".", apFindData->cFileName) != 0)
                                && (::wcscmp(L"..", apFindData->cFileName) != 0)) {
                            // Create full path to file/dir by appending the file name
                            // to the current path. Add it to buffer.
                            std::wstring subPath = p_Path + L"\\" + apFindData->cFileName;
                            apSubPaths->push_back(subPath);
                        }
                    } while (::FindNextFileW(hFind, apFindData.get()));

                    // Now iterate buffer and recurse.
                    FilesV::const_iterator it, end = apSubPaths->end();
                    for (it = apSubPaths->begin(); it != end; ++it) {
                        ModifyReadOnlyFlag(*it, p_FlagState);
                    }
                } catch (...) {
                    ::FindClose(hFind);
                    throw;
                }
                ::FindClose(hFind);
            }
        }
    }
}
